#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 10.使用PyTorch完成Fashion-MNIST分类.py
# @Author: Richard Chiming Xu
# @Date  : 2021/11/16
# @Desc  :

import torch
from torch import nn
from torchvision import transforms
import torchvision
from torch.utils.data import Dataset

# 判定是否使用GPU计算
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

config = {
    'BATCH_SIZE': 128,
    'LEARNING_RATE': 0.01,
    'EPOCHS': 50
}


# 初始化图片转换
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 读取数据集
trainset = torchvision.datasets.FashionMNIST('D:/env/sample_data/', download=True, train=True, transform=transform)
testset = torchvision.datasets.FashionMNIST('D:/env/sample_data/', download=True, train=False, transform=transform)

# 转换dataloader
trainloader = torch.utils.data.DataLoader(trainset, batch_size=config['BATCH_SIZE'], shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=config['BATCH_SIZE'], shuffle=True)

# 搭建网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU())  # 16, 28, 28
        self.pool1 = nn.MaxPool2d(2)  # 16, 14, 14
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=3),
            nn.BatchNorm2d(32),
            nn.ReLU())  # 32, 12, 12
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU())  # 64, 10, 10
        self.layer4 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3),
            nn.BatchNorm2d(128),
            nn.ReLU())  # 128, 8, 8
        self.pool2 = nn.MaxPool2d(2)  # 128, 4, 4
        self.fc1 = nn.Linear(4 * 4 * 128, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        out = self.layer1(x)
        out = self.pool1(out)

        out = self.layer2(out)

        out = self.layer3(out)
        out = self.layer4(out)
        out = self.pool2(out)

        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.fc2(out)

        return out

cnn =CNN()
cnn.cuda(device)
print(cnn)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.Adam(cnn.parameters(), lr = config['LEARNING_RATE'])


losses = []
for epoch in range(config['EPOCHS']):
    for i, (images, labels) in enumerate(trainloader):
        images = images.float().to(device)
        labels = labels.to(device)


        optimizer.zero_grad()  # 梯度清零
        outputs = cnn(images)
        loss = criterion(outputs, labels)  # 计算损失函数
        loss.backward()
        optimizer.step()
        losses.append(loss.cpu().data.item())
        if (i + 1) % 100 == 0:
            print('Epoch : %d/%d, Iter : %d/%d,  Loss: %.4f' % (
            epoch + 1, config['EPOCHS'], i + 1, len(trainloader) // config['BATCH_SIZE'], loss.data.item()))
